package regressor

import (
	"errors"
	"fmt"
	"math"
	"math/rand"
	"perceptron/perceptron"
	"perceptron/perceptron/functions"

	"gonum.org/v1/gonum/blas/blas64"
)

type Regressor struct {
	layers         []perceptron.Layer
	inputParamSize int
	iters          int
	learningRate   float64

	activation    func(float64) float64 // activation function
	dActivation   func(float64) float64 // derivative of activation function
	lastLayerFunc func(float64) float64 // last layer activation function
	lasLayerDFunc func(float64) float64 // last layer derivative function
}

func (r *Regressor) Init(layers []uint, iters int, inputParamSize int, learningRate float64) error {
	if iters <= 0 || inputParamSize <= 0 || layers == nil || len(layers) == 0 || learningRate < 0 {
		return errors.New("invalid params")
	}
	r.layers = make([]perceptron.Layer, 0)
	r.inputParamSize = inputParamSize
	r.iters = iters
	r.learningRate = learningRate
	r.activation = functions.ReLU
	r.dActivation = functions.ReLuDerivative
	r.lastLayerFunc = functions.ReLU
	r.lasLayerDFunc = functions.ReLuDerivative

	layers = append([]uint{uint(inputParamSize)}, layers...)
	layers = append(layers, 1)

	// Инициализация слоев и их нейронов
	for iLayer, size := range layers {
		var neurons []perceptron.Neuron
		for i := uint(0); i < size; i++ {
			// Инициализация весов случайными значениями
			if iLayer == 0 {
				neurons = append(neurons, perceptron.Neuron{Weights: nil})
			} else {
				weights := make([]float64, layers[iLayer-1])
				for j := range weights {
					weights[j] = rand.Float64()
				}
				neurons = append(neurons, perceptron.Neuron{Weights: weights})
			}
		}
		r.layers = append(r.layers, perceptron.Layer{
			Neurons: neurons,
		})
	}

	fmt.Println("слои")
	for i, l := range r.layers {
		fmt.Println("слой", i, "нейронов", len(l.Neurons))
	}

	return nil
}

func (r *Regressor) Fit(x [][]float64, y []float64) {
	for i := 0; i < r.iters; i++ {
		//прямой проход
		r.Run(x[i%len(x)])

		//обратное распространение
		r.backPropagation([]float64{y[i%len(y)]})

	}

}

func (r *Regressor) Run(inputs []float64) []float64 {
	for i, layer := range r.layers {
		if layer.Neurons[0].Weights == nil {
			for i, neuron := range layer.Neurons {
				neuron.Value = inputs[i]
			}
			continue
		}

		outputs := make([]float64, 0)

		inputsVector := blas64.Vector{
			N:    len(inputs),
			Data: inputs,
			Inc:  1,
		}
		for in, neuron := range layer.Neurons {
			weightsVector := blas64.Vector{
				N:    len(neuron.Weights),
				Data: neuron.Weights,
				Inc:  1,
			}

			dotProduct := blas64.Dot(inputsVector, weightsVector)

			// Применение функции активации
			if i != len(r.layers)-1 {
				layer.Neurons[in].Value = r.activation(dotProduct)
			} else {
				layer.Neurons[in].Value = r.lastLayerFunc(dotProduct)
			}

			if math.IsNaN(layer.Neurons[in].Value) {
				panic("NAN VALUE")
			}

			outputs = append(outputs, layer.Neurons[in].Value)
		}
		inputs = outputs

	}
	return inputs
}

// обратное расспространение ошибки
func (r *Regressor) backPropagation(y []float64) {
	//вычисляем ошибки
	errs := make([]float64, len(y))
	for i := range y {
		errs[i] = r.layers[len(r.layers)-1].Neurons[i].Value - y[i]
	}

	// Проходимся по слоям в обратном порядке
	for i := len(r.layers) - 1; i > 0; i-- {
		// Если текущий слой - выходной, вычисляем градиенты на основе ошибок
		if i == len(r.layers)-1 {
			errs = r.backPropOneLayer(errs, i, r.lasLayerDFunc)
		} else {
			errs = r.backPropOneLayer(errs, i, r.dActivation)
		}
	}

}

func (r *Regressor) backPropOneLayer(errs []float64, curLayer int, df func(float64) float64) []float64 {
	newErrs := make([]float64, len(r.layers[curLayer-1].Neurons))
	for in, neuron := range r.layers[curLayer].Neurons {
		// Вычисляем градиент для каждого нейрона на выходном слое
		gradient := 0.0

		gradient += errs[in] * df(neuron.Value)

		// Обновляем веса нейрона на основе градиента и скорости обучения
		for k := range neuron.Weights {
			neuron.Weights[k] -= r.learningRate * gradient * r.layers[curLayer-1].Neurons[k].Value

			if math.IsNaN(neuron.Weights[k]) {
				panic("NAN in neuron weight")
			}
			newErrs[k] += neuron.Weights[k] * gradient

		}
		//обновляем ошибки для следующего слоя
	}
	return newErrs
}

func (r *Regressor) Test(x [][]float64, y []float64) (string, float64) {
	err := 0.
	for i, inp := range x {
		out := r.Run(inp)
		err += math.Pow(y[i]-out[0], 2)
	}
	return "Сумма квадратов ошибок", err
}
